#!/bin/bash

# Qwen3-14B 基础环境部署脚本 (适用于 Ubuntu 22.04 + GPU，FP16 模式)
set -e  # 遇到错误立即退出

echo "开始安装Qwen3-14B基础环境..."

# 1. 更新系统并安装基础依赖
echo "步骤1/6: 更新系统包及安装基础工具..."
sudo apt update && sudo apt upgrade -y
sudo apt install -y git-lfs wget curl build-essential ca-certificates

# 2. 安装Miniconda (Python环境管理)
echo "步骤2/6: 安装Miniconda..."
if ! command -v conda &> /dev/null; then
    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
    bash miniconda.sh -b -p $HOME/miniconda
    rm miniconda.sh

    # === 关键修改：立即将Miniconda的路径添加到当前Shell的PATH中 ===
    export PATH="$HOME/miniconda/bin:$PATH"
    # 同时也写入.bashrc以供将来使用
    echo 'export PATH="$HOME/miniconda/bin:$PATH"' >> ~/.bashrc

    # 初始化 conda（支持 conda activate）
    echo "初始化 Conda..."
    conda init bash
    # 对于 zsh 用户也初始化
    if [ -n "$ZSH_VERSION" ]; then
        conda init zsh
    fi
    
    echo "注意: 如果当前 shell 无法使用 'conda activate'，请运行:"
    echo "  source ~/.bashrc  (对于 bash)"
    echo "  或"
    echo "  source ~/.zshrc   (对于 zsh)"
    echo "  或者重新打开终端"

else
    echo "Conda 已安装，跳过。"
    # 即使 conda 已安装，也检查是否已初始化
    if ! grep -q "conda initialize" ~/.bashrc 2>/dev/null; then
        echo "检测到 Conda 未初始化，正在初始化..."
        conda init bash
        if [ -n "$ZSH_VERSION" ]; then
            conda init zsh
        fi
        echo "请运行: source ~/.bashrc 或重新打开终端"
    fi
fi

# 3. 创建并激活Conda虚拟环境 (Python 3.10)
echo "步骤3/6: 创建Conda虚拟环境..."
# 现在conda命令在当前脚本中应该可用了
conda create -n qwen314b python=3.10 -y

# 激活环境（重要！）
echo "激活 qwen314b 环境..."
# 初始化 conda（如果还没有）
if [ -f "$HOME/miniconda/etc/profile.d/conda.sh" ]; then
    source "$HOME/miniconda/etc/profile.d/conda.sh"
elif [ -f "$HOME/anaconda3/etc/profile.d/conda.sh" ]; then
    source "$HOME/anaconda3/etc/profile.d/conda.sh"
fi

# 激活环境
conda activate qwen314b

# 验证环境
echo "验证 Python 版本..."
python --version || python3 --version

# 4. 配置PyPI镜像加速及升级pip
echo "步骤4/6: 配置Python环境..."
# 检测可用的 Python 命令（某些系统只有 python3）
if command -v python3 &> /dev/null; then
    PYTHON_CMD="python3"
elif command -v python &> /dev/null; then
    PYTHON_CMD="python"
else
    echo "错误: 未找到 python 或 python3 命令"
    exit 1
fi
echo "使用 Python 命令: $PYTHON_CMD"
echo "Python 版本: $($PYTHON_CMD --version)"
echo "Python 路径: $(which $PYTHON_CMD)"

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
$PYTHON_CMD -m pip install --upgrade pip

# 5. 安装PyTorch (CUDA 12.1) 及核心库
echo "步骤5/6: 安装PyTorch和vLLM..."
# 注意: torchaudio 在某些 Python 版本可能不可用，如果失败可以跳过
echo "安装 PyTorch (torch, torchvision, torchaudio)..."
if $PYTHON_CMD -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121; then
    echo "✓ PyTorch 安装成功"
else
    echo "⚠ torchaudio 安装失败，尝试只安装 torch 和 torchvision..."
    $PYTHON_CMD -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
    echo "注意: torchaudio 未安装，但这通常不影响 vLLM 的使用"
fi

echo "安装 vLLM..."
$PYTHON_CMD -m pip install vllm

# 6. 创建模型目录并下载Qwen3-14B模型 (通过ModelScope)
echo "步骤6/6: 创建模型目录并下载Qwen3-14B模型..."
# 创建模型存储目录
MODEL_BASE_DIR="/data/models"
# ModelScope 的 snapshot_download('Qwen/Qwen3-14B', cache_dir='...') 会在 cache_dir 下创建 Qwen/Qwen3-14B 目录
# 如果 cache_dir=/data/models/Qwen，实际路径会是 /data/models/Qwen/Qwen/Qwen3-14B
MODEL_CACHE_DIR="$MODEL_BASE_DIR/Qwen"
MODEL_FINAL_DIR="$MODEL_BASE_DIR/Qwen/Qwen/Qwen3-14B"
echo "创建模型存储目录: $MODEL_CACHE_DIR"
# 如果 /data 目录不存在，先创建它
if [ ! -d "/data" ]; then
    echo "创建 /data 目录..."
    sudo mkdir -p /data
    sudo chown $USER:$USER /data
    sudo chmod 755 /data
fi
# 创建模型目录
sudo mkdir -p "$MODEL_CACHE_DIR"
# 设置目录权限，确保当前用户有读写权限
sudo chown -R $USER:$USER "$MODEL_BASE_DIR"
sudo chmod -R 755 "$MODEL_BASE_DIR"
echo "✓ 模型目录已创建并设置权限: $MODEL_CACHE_DIR"
echo "  模型将下载至: $MODEL_FINAL_DIR"

# 下载模型
# 注意: ModelScope 的 snapshot_download 会在 cache_dir 下创建模型名称对应的目录
# 对于 'Qwen/Qwen3-14B'，会在 cache_dir 下创建 Qwen/Qwen3-14B 目录
# 所以如果 cache_dir=/data/models/Qwen，最终路径是 /data/models/Qwen/Qwen/Qwen3-14B
echo "安装 ModelScope..."
$PYTHON_CMD -m pip install modelscope
$PYTHON_CMD -c "
from modelscope import snapshot_download
import os
model_dir = snapshot_download('Qwen/Qwen3-14B', cache_dir='$MODEL_CACHE_DIR')
print(f'模型已下载至: {model_dir}')
# 验证实际路径
if os.path.exists('$MODEL_FINAL_DIR'):
    print(f'✓ 模型文件位于: $MODEL_FINAL_DIR')
else:
    # 如果路径不对，尝试查找实际路径
    import glob
    possible_paths = glob.glob('$MODEL_BASE_DIR/Qwen*/Qwen*/Qwen3-14B')
    if possible_paths:
        actual_path = possible_paths[0]
        print(f'⚠ 检测到模型实际路径: {actual_path}')
        print(f'  请确认此路径是否正确')
    else:
        print(f'⚠ 未找到模型文件，请检查下载是否成功')
"

echo "基础环境部署完成！"
echo ""
echo "========================================="
echo "下一步：激活环境"
echo "========================================="
echo "如果遇到 'conda activate' 错误，请运行修复脚本："
echo "  bash $(dirname "$0")/fix_conda_init.sh"
echo ""
echo "或者手动运行："
echo "  source ~/.bashrc  (bash 用户)"
echo "  或"
echo "  source ~/.zshrc   (zsh 用户)"
echo "  或者重新打开终端"
echo ""
echo "然后激活环境："
echo "  conda activate qwen314b"
echo ""
echo "========================================="
echo "测试环境"
echo "========================================="
echo "1. 运行测试脚本验证安装："
echo "   bash $(dirname "$0")/test_qwen314b.sh"
echo "   或"
echo "   conda activate qwen314b"
echo "   python3 $(dirname "$0")/quick_test_qwen.py  # 或使用 python"
echo ""
echo "2. 启动 vLLM API 服务（可选）："
echo "   # 注意: 某些系统只有 python3，使用 python3 或 python"
echo "   python3 -m vllm.entrypoints.openai.api_server \\"
echo "     --model /data/models/Qwen/Qwen/Qwen3-14B \\"
echo "     --trust-remote-code"
echo ""
echo "注意: 模型实际路径为 /data/models/Qwen/Qwen/Qwen3-14B"
echo "========================================="

